import pygame
import random
import json
import os
from collections import deque

# 初始化pygame
pygame.init()

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
GRAY = (200, 200, 200)

# 游戏设置
GRID_SIZE = 20
GRID_WIDTH = 30
GRID_HEIGHT = 20
SCREEN_WIDTH = GRID_SIZE * GRID_WIDTH
SCREEN_HEIGHT = GRID_SIZE * GRID_HEIGHT + 200  # 增加UI区域高度
GAME_AREA_HEIGHT = GRID_SIZE * GRID_HEIGHT

# 创建游戏窗口
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("自动贪吃蛇")

# 游戏时钟
clock = pygame.time.Clock()

# 排行榜文件
HIGH_SCORES_FILE = "high_scores.json"


class SnakeGame:
    def __init__(self):
        self.reset_game()
        self.load_high_scores()

    def reset_game(self):
        # 蛇的初始位置和长度
        self.snake = deque([(GRID_WIDTH // 2, GRID_HEIGHT // 2)])
        self.direction = (1, 0)  # 初始向右移动
        self.next_direction = (1, 0)

        # 生成食物
        self.food = self.generate_food()

        # 游戏状态
        self.score = 0
        self.game_over = False
        self.game_speed = 20  # 默认游戏速度

    def generate_food(self):
        while True:
            food = (random.randint(0, GRID_WIDTH - 1), random.randint(0, GRID_HEIGHT - 1))
            if food not in self.snake:
                return food

    def load_high_scores(self):
        self.high_scores = []
        if os.path.exists(HIGH_SCORES_FILE):
            with open(HIGH_SCORES_FILE, 'r') as f:
                self.high_scores = json.load(f)

    def save_high_scores(self):
        with open(HIGH_SCORES_FILE, 'w') as f:
            json.dump(self.high_scores, f)

    def add_high_score(self, score):
        self.high_scores.append(score)
        self.high_scores.sort(reverse=True)
        self.high_scores = self.high_scores[:5]  # 只保留前5条记录
        self.save_high_scores()

    def auto_move(self):
        # A*寻路算法实现
        head_x, head_y = self.snake[0]
        food_x, food_y = self.food

        # 启发式函数：曼哈顿距离
        def heuristic(x, y):
            return abs(x - food_x) + abs(y - food_y)

        # 获取可行方向
        def get_valid_directions(x, y, prev_dir=None):
            directions = []
            for dx, dy in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
                # 不能反向移动
                if prev_dir and (dx == -prev_dir[0] and dy == -prev_dir[1]):
                    continue

                nx = x + dx
                ny = y + dy

                # 检查边界和蛇身碰撞
                if (0 <= nx < GRID_WIDTH and
                        0 <= ny < GRID_HEIGHT and
                        (nx, ny) not in self.snake):
                    directions.append((dx, dy))
            return directions

        # A*算法实现
        open_set = [(head_x, head_y)]
        came_from = {}
        g_score = {(head_x, head_y): 0}
        f_score = {(head_x, head_y): heuristic(head_x, head_y)}

        while open_set:
            current = min(open_set, key=lambda pos: f_score[pos])

            if current == (food_x, food_y):
                # 重建路径
                path = []
                while current in came_from:
                    path.append(current)
                    current = came_from[current]
                if path:
                    next_pos = path[-1]
                    self.next_direction = (next_pos[0] - head_x, next_pos[1] - head_y)
                return

            open_set.remove(current)

            for dx, dy in get_valid_directions(current[0], current[1]):
                neighbor = (current[0] + dx, current[1] + dy)

                # 计算临时g分数
                tentative_g_score = g_score[current] + 1

                if neighbor not in g_score or tentative_g_score < g_score[neighbor]:
                    came_from[neighbor] = current
                    g_score[neighbor] = tentative_g_score
                    f_score[neighbor] = g_score[neighbor] + heuristic(neighbor[0], neighbor[1])
                    if neighbor not in open_set:
                        open_set.append(neighbor)

        # 如果没有找到路径，使用安全区域评估
        best_dir = None
        best_score = -float('inf')

        for dx, dy in get_valid_directions(head_x, head_y, self.direction):
            # 计算安全区域评分
            safe_area = 0
            for ddx, ddy in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
                nx = head_x + dx + ddx
                ny = head_y + dy + ddy
                if (0 <= nx < GRID_WIDTH and
                        0 <= ny < GRID_HEIGHT and
                        (nx, ny) not in self.snake):
                    safe_area += 1

            # 检查是否为死胡同
            is_dead_end = True
            visited = set()
            queue = deque([(head_x + dx, head_y + dy)])

            while queue:
                x, y = queue.popleft()
                if (x, y) in visited or (x, y) in self.snake or \
                        not (0 <= x < GRID_WIDTH and 0 <= y < GRID_HEIGHT):
                    continue

                visited.add((x, y))

                # 如果找到多个出口，不是死胡同
                if len(visited) > 2:
                    is_dead_end = False
                    break

                # 向四个方向扩展
                for ddx, ddy in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
                    queue.append((x + ddx, y + ddy))

            # 当蛇身长度超过75时，优先考虑空白区域
            if len(self.snake) > 75:
                # 计算当前方向可到达的空白区域大小
                blank_area = 0
                visited = set()
                queue = deque([(head_x + dx, head_y + dy)])

                while queue:
                    x, y = queue.popleft()
                    if (x, y) in visited or (x, y) in self.snake or \
                            not (0 <= x < GRID_WIDTH and 0 <= y < GRID_HEIGHT):
                        continue

                    visited.add((x, y))
                    blank_area += 1

                    # 向四个方向扩展
                    for ddx, ddy in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
                        queue.append((x + ddx, y + ddy))

                # 如果食物不在边缘，避免靠近边缘
                edge_penalty = 0
                if not (food_x == 0 or food_x == GRID_WIDTH - 1 or
                        food_y == 0 or food_y == GRID_HEIGHT - 1):
                    if (head_x + dx == 0 or head_x + dx == GRID_WIDTH - 1 or
                            head_y + dy == 0 or head_y + dy == GRID_HEIGHT - 1):
                        edge_penalty = 50

                # 评分 = 空白区域 - 边缘惩罚 - 死胡同惩罚
                # 增加死胡同惩罚权重，并考虑蛇身长度
                dead_end_penalty = 150 + len(self.snake) * 2 if is_dead_end else 0
                score = blank_area - edge_penalty - dead_end_penalty
            else:
                # 评分 = 安全区域 - 距离食物距离 - 死胡同惩罚
                # 增加死胡同惩罚权重，并考虑蛇身长度
                dead_end_penalty = 150 + len(self.snake) * 2 if is_dead_end else 0
                score = safe_area - heuristic(head_x + dx, head_y + dy) - dead_end_penalty

            if score > best_score:
                best_score = score
                best_dir = (dx, dy)

        if best_dir:
            self.next_direction = best_dir
        else:
            # 没有可行方向，随机选择一个方向
            self.next_direction = self.direction

    def update(self):
        if self.game_over:
            return

        # 自动寻路
        self.auto_move()
        self.direction = self.next_direction

        # 移动蛇
        head_x, head_y = self.snake[0]
        new_x = head_x + self.direction[0]
        new_y = head_y + self.direction[1]

        # 检查碰撞
        if (new_x < 0 or new_x >= GRID_WIDTH or
                new_y < 0 or new_y >= GRID_HEIGHT or
                (new_x, new_y) in self.snake):
            self.game_over = True
            self.add_high_score(self.score)
            return

        # 移动蛇
        self.snake.appendleft((new_x, new_y))

        # 检查是否吃到食物
        if (new_x, new_y) == self.food:
            self.score += 1
            self.food = self.generate_food()
            self.steps_without_food = 0
        else:
            self.snake.pop()

    def draw(self):
        # 绘制游戏区域背景
        screen.fill(BLACK)
        pygame.draw.rect(screen, WHITE, (0, 0, SCREEN_WIDTH, GAME_AREA_HEIGHT))

        # 绘制蛇
        for i, segment in enumerate(self.snake):
            x, y = segment
            # 蛇尾最后5格使用渐变颜色
            if i >= len(self.snake) - 5:
                # 计算渐变比例 (0.2到1.0)
                ratio = 0.2 + 0.8 * (len(self.snake) - i) / 5
                color = (0, int(255 * ratio), 0)
            else:
                color = GREEN
            pygame.draw.rect(screen, color,
                             (x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE, GRID_SIZE))

        # 绘制蛇头
        head_x, head_y = self.snake[0]
        pygame.draw.rect(screen, BLUE,
                         (head_x * GRID_SIZE, head_y * GRID_SIZE, GRID_SIZE, GRID_SIZE))

        # 绘制食物
        food_x, food_y = self.food
        pygame.draw.rect(screen, RED,
                         (food_x * GRID_SIZE, food_y * GRID_SIZE, GRID_SIZE, GRID_SIZE))

        # 绘制UI区域
        pygame.draw.rect(screen, GRAY, (0, GAME_AREA_HEIGHT, SCREEN_WIDTH, 200))

        # 显示分数
        font = pygame.font.SysFont(None, 30)
        score_text = font.render(f"Score: {self.score}", True, BLACK)
        screen.blit(score_text, (10, GAME_AREA_HEIGHT + 10))

        # 显示重新开始按钮
        restart_button = pygame.Rect(SCREEN_WIDTH - 120, GAME_AREA_HEIGHT + 10, 100, 30)
        pygame.draw.rect(screen, BLUE, restart_button)
        restart_text = font.render("Restart", True, WHITE)
        screen.blit(restart_text, (SCREEN_WIDTH - 110, GAME_AREA_HEIGHT + 15))

        # 绘制速度控制按钮
        pygame.draw.rect(screen, GREEN, (SCREEN_WIDTH - 240, GAME_AREA_HEIGHT + 10, 50, 30))
        speed_up_text = font.render("+", True, WHITE)
        screen.blit(speed_up_text, (SCREEN_WIDTH - 230, GAME_AREA_HEIGHT + 15))

        pygame.draw.rect(screen, RED, (SCREEN_WIDTH - 180, GAME_AREA_HEIGHT + 10, 50, 30))
        speed_down_text = font.render("-", True, WHITE)
        screen.blit(speed_down_text, (SCREEN_WIDTH - 170, GAME_AREA_HEIGHT + 15))

        # 显示排行榜
        high_scores_text = font.render("High Scores:", True, WHITE)
        screen.blit(high_scores_text, (10, GAME_AREA_HEIGHT + 50))

        for i, score in enumerate(self.high_scores[:5]):
            score_text = font.render(f"{i + 1}. {score}", True, WHITE)
            screen.blit(score_text, (10, GAME_AREA_HEIGHT + 80 + i * 40))

        # 游戏结束显示
        if self.game_over:
            font = pygame.font.SysFont(None, 50)
            game_over_text = font.render("Game Over!", True, RED)
            screen.blit(game_over_text,
                        (SCREEN_WIDTH // 2 - 100, GAME_AREA_HEIGHT // 2 - 25))


# 主游戏循环
def main():
    game = SnakeGame()
    running = True

    # 定义速度控制按钮
    speed_up_button = pygame.Rect(SCREEN_WIDTH - 240, GAME_AREA_HEIGHT + 10, 50, 30)
    speed_down_button = pygame.Rect(SCREEN_WIDTH - 180, GAME_AREA_HEIGHT + 10, 50, 30)

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                # 检查是否点击了重新开始按钮
                mouse_x, mouse_y = pygame.mouse.get_pos()
                restart_button = pygame.Rect(SCREEN_WIDTH - 120, GAME_AREA_HEIGHT + 10, 100, 30)
                if restart_button.collidepoint(mouse_x, mouse_y):
                    game.reset_game()
                elif speed_up_button.collidepoint(mouse_x, mouse_y):
                    game.game_speed = min(100, game.game_speed + 10)
                elif speed_down_button.collidepoint(mouse_x, mouse_y):
                    game.game_speed = max(2, game.game_speed - 2)

        game.update()
        game.draw()

        pygame.display.flip()
        clock.tick(game.game_speed)  # 使用变量控制游戏速度

    pygame.quit()


if __name__ == "__main__":
    main()

